use crate::{
    ast::{Expression, Type},
    error::{CompilerError, Result},
    llvm_backend::LLVMBackend,
};
use inkwell::{
    types::BasicType,
    values::BasicValueEnum,
    AddressSpace,
    IntPredicate,
};

impl<'ctx> LLVMBackend<'ctx> {

    /// Generate string.join(vector) method call
    /// Joins a vector of strings with the separator string
    ///
    /// Example: ", ".join(vec) where vec is Vector<String>
    /// Returns the joined string
    pub(in crate::llvm_backend) fn generate_string_join_method(
        &mut self,
        separator_var_name: &str,
        arguments: &[Expression],
    ) -> Result<BasicValueEnum<'ctx>> {
        // Validate arguments
        if arguments.len() != 1 {
            return Err(CompilerError::internal(
                "join() requires exactly one argument (Vector<String>)",
            ));
        }

        // Get the separator string pointer
        let separator_ptr = if let Some(var_value) = self.variable_map.get(separator_var_name) {
            var_value.into_pointer_value()
        } else {
            return Err(CompilerError::internal(&format!(
                "Undefined separator variable: {}",
                separator_var_name
            )));
        };

        // Load the separator string pointer value (i8*)
        let i8_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
        let separator_value =
            self.builder
                .build_load(i8_ptr_type, separator_ptr, "separator_val")?;
        let separator_str_ptr = separator_value.into_pointer_value();

        // Get the vector argument
        let vector_arg = &arguments[0];

        // Get the vector variable name and value
        let (vector_var_name, vector_ptr) = if let Expression::Variable(var_name) = vector_arg {
            let vec_ptr = if let Some(var_value) = self.variable_map.get(var_name) {
                var_value.into_pointer_value()
            } else {
                return Err(CompilerError::internal(&format!(
                    "Undefined vector variable: {}",
                    var_name
                )));
            };
            (var_name.clone(), vec_ptr)
        } else {
            return Err(CompilerError::internal(
                "join() requires a vector variable as argument",
            ));
        };

        // Get the vector type and verify it's Vector<String>
        let element_type =
            if let Some(Type::Vec(vec_type)) = self.variable_type_map.get(&vector_var_name) {
                if !matches!(
                    *vec_type.element_type,
                    Type::Basic(crate::ast::types::BasicType::String)
                ) {
                    return Err(CompilerError::internal(
                        "join() requires Vector<String> as argument",
                    ));
                }
                *vec_type.element_type.clone()
            } else {
                return Err(CompilerError::internal("join() requires a vector variable"));
            };

        // Get element LLVM type (i8* for String)
        let element_llvm_type = self.nrc_type_to_llvm_type(&Some(element_type))?;

        // Create slice struct type: {i8**, i64, i64} for Vector<String>
        let ptr_type = element_llvm_type.ptr_type(AddressSpace::default());
        let slice_struct_type = self.context.struct_type(
            &[
                ptr_type.into(),
                self.context.i64_type().into(),
                self.context.i64_type().into(),
            ],
            false,
        );

        // Load the vector struct
        let vector_struct =
            self.builder
                .build_load(slice_struct_type, vector_ptr, "vector_struct")?;
        if !vector_struct.is_struct_value() {
            return Err(CompilerError::internal("Vector value is not a struct"));
        }
        let vec_struct = vector_struct.into_struct_value();

        // Extract fields: ptr, len, cap
        let strings_ptr = self
            .builder
            .build_extract_value(vec_struct, 0, "strings_ptr")?
            .into_pointer_value();
        let vec_len = self
            .builder
            .build_extract_value(vec_struct, 1, "vec_len")?
            .into_int_value();

        // Get strlen function for calculating string lengths
        let i32_type = self.context.i32_type();
        let i64_type = self.context.i64_type();
        let i8_type = self.context.i8_type();

        let strlen_fn_type = i32_type.fn_type(&[i8_ptr_type.into()], false);
        let strlen_fn = if let Some(func) = self.module.get_function("strlen") {
            func
        } else {
            self.module.add_function("strlen", strlen_fn_type, None)
        };

        // Create basic blocks
        let current_func = self.current_function.unwrap();
        let empty_check_block = self
            .context
            .append_basic_block(current_func, "join_empty_check");
        let single_check_block = self
            .context
            .append_basic_block(current_func, "join_single_check");
        let calc_len_loop = self
            .context
            .append_basic_block(current_func, "join_calc_len_loop");
        let calc_len_body = self
            .context
            .append_basic_block(current_func, "join_calc_len_body");
        let calc_len_end = self
            .context
            .append_basic_block(current_func, "join_calc_len_end");
        let alloc_block = self.context.append_basic_block(current_func, "join_alloc");
        let copy_loop = self
            .context
            .append_basic_block(current_func, "join_copy_loop");
        let copy_body = self
            .context
            .append_basic_block(current_func, "join_copy_body");
        let return_block = self.context.append_basic_block(current_func, "join_return");

        // Jump to empty check
        self.builder.build_unconditional_branch(empty_check_block)?;

        // Empty check: if len == 0, return empty string
        self.builder.position_at_end(empty_check_block);
        let is_empty = self.builder.build_int_compare(
            IntPredicate::EQ,
            vec_len,
            i64_type.const_int(0, false),
            "is_empty",
        )?;
        self.builder
            .build_conditional_branch(is_empty, return_block, single_check_block)?;

        // Single check: if len == 1, return the first string
        self.builder.position_at_end(single_check_block);
        let is_single = self.builder.build_int_compare(
            IntPredicate::EQ,
            vec_len,
            i64_type.const_int(1, false),
            "is_single",
        )?;

        // Get first string for single element case
        let first_str_ptr_loc = unsafe {
            self.builder.build_gep(
                element_llvm_type,
                strings_ptr,
                &[i64_type.const_int(0, false)],
                "first_str_ptr_loc",
            )?
        };
        let first_str = self
            .builder
            .build_load(element_llvm_type, first_str_ptr_loc, "first_str")?
            .into_pointer_value();

        self.builder
            .build_conditional_branch(is_single, return_block, calc_len_loop)?;

        // Calculate total length needed
        self.builder.position_at_end(calc_len_loop);

        // Allocate variables for loop
        let total_len_alloca = self.builder.build_alloca(i64_type, "total_len")?;
        let idx_alloca = self.builder.build_alloca(i64_type, "idx")?;
        self.builder
            .build_store(total_len_alloca, i64_type.const_int(0, false))?;
        self.builder
            .build_store(idx_alloca, i64_type.const_int(0, false))?;

        // Calculate separator length once
        let sep_len_i32 = self
            .builder
            .build_call(strlen_fn, &[separator_str_ptr.into()], "sep_len_i32")?
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_int_value();
        let sep_len = self
            .builder
            .build_int_z_extend(sep_len_i32, i64_type, "sep_len")?;

        self.builder.build_unconditional_branch(calc_len_body)?;

        // Calculate length loop body
        self.builder.position_at_end(calc_len_body);
        let idx = self
            .builder
            .build_load(i64_type, idx_alloca, "idx")?
            .into_int_value();

        // Check if idx < vec_len
        let continue_calc = self.builder.build_int_compare(
            IntPredicate::SLT,
            idx,
            vec_len,
            "continue_calc",
        )?;
        self.builder
            .build_conditional_branch(continue_calc, calc_len_end, alloc_block)?;

        // Calculate length body continuation
        self.builder.position_at_end(calc_len_end);

        // Get current string pointer
        let str_ptr_loc = unsafe {
            self.builder
                .build_gep(element_llvm_type, strings_ptr, &[idx], "str_ptr_loc")?
        };
        let current_str = self
            .builder
            .build_load(element_llvm_type, str_ptr_loc, "current_str")?
            .into_pointer_value();

        // Get string length
        let str_len_i32 = self
            .builder
            .build_call(strlen_fn, &[current_str.into()], "str_len_i32")?
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_int_value();
        let str_len = self
            .builder
            .build_int_z_extend(str_len_i32, i64_type, "str_len")?;

        // Add to total
        let total_len = self
            .builder
            .build_load(i64_type, total_len_alloca, "total_len")?
            .into_int_value();
        let new_total = self
            .builder
            .build_int_add(total_len, str_len, "new_total")?;
        self.builder.build_store(total_len_alloca, new_total)?;

        // Increment idx
        let next_idx = self
            .builder
            .build_int_add(idx, i64_type.const_int(1, false), "next_idx")?;
        self.builder.build_store(idx_alloca, next_idx)?;

        self.builder.build_unconditional_branch(calc_len_body)?;

        // Allocate memory block
        self.builder.position_at_end(alloc_block);
        let total_len_final = self
            .builder
            .build_load(i64_type, total_len_alloca, "total_len_final")?
            .into_int_value();

        // Add separator lengths: sep_len * (vec_len - 1)
        let vec_len_minus_one = self.builder.build_int_sub(
            vec_len,
            i64_type.const_int(1, false),
            "vec_len_minus_one",
        )?;
        let total_sep_len =
            self.builder
                .build_int_mul(sep_len, vec_len_minus_one, "total_sep_len")?;
        let total_with_sep =
            self.builder
                .build_int_add(total_len_final, total_sep_len, "total_with_sep")?;

        // Add 1 for null terminator
        let total_with_null = self.builder.build_int_add(
            total_with_sep,
            i64_type.const_int(1, false),
            "total_with_null",
        )?;

        // Allocate memory using malloc
        let malloc_fn_type = i8_ptr_type.fn_type(&[i64_type.into()], false);
        let malloc_fn = if let Some(func) = self.module.get_function("malloc") {
            func
        } else {
            self.module.add_function("malloc", malloc_fn_type, None)
        };

        let result_str = self
            .builder
            .build_call(malloc_fn, &[total_with_null.into()], "result_str")?
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // Initialize strcpy and strcat functions
        let strcpy_fn_type = i8_ptr_type.fn_type(&[i8_ptr_type.into(), i8_ptr_type.into()], false);
        let strcpy_fn = if let Some(func) = self.module.get_function("strcpy") {
            func
        } else {
            self.module.add_function("strcpy", strcpy_fn_type, None)
        };

        let strcat_fn_type = i8_ptr_type.fn_type(&[i8_ptr_type.into(), i8_ptr_type.into()], false);
        let strcat_fn = if let Some(func) = self.module.get_function("strcat") {
            func
        } else {
            self.module.add_function("strcat", strcat_fn_type, None)
        };

        // Initialize result with empty string (set first byte to 0)
        self.builder
            .build_store(result_str, i8_type.const_int(0, false))?;

        // Reset idx for copy loop
        self.builder
            .build_store(idx_alloca, i64_type.const_int(0, false))?;
        self.builder.build_unconditional_branch(copy_loop)?;

        // Copy loop
        self.builder.position_at_end(copy_loop);
        let copy_idx = self
            .builder
            .build_load(i64_type, idx_alloca, "copy_idx")?
            .into_int_value();

        // Check if copy_idx < vec_len
        let continue_copy = self.builder.build_int_compare(
            IntPredicate::SLT,
            copy_idx,
            vec_len,
            "continue_copy",
        )?;
        self.builder
            .build_conditional_branch(continue_copy, copy_body, return_block)?;

        // Copy body
        self.builder.position_at_end(copy_body);

        // Get current string
        let copy_str_ptr_loc = unsafe {
            self.builder.build_gep(
                element_llvm_type,
                strings_ptr,
                &[copy_idx],
                "copy_str_ptr_loc",
            )?
        };
        let copy_str = self
            .builder
            .build_load(element_llvm_type, copy_str_ptr_loc, "copy_str")?
            .into_pointer_value();

        // Check if this is the first string
        let is_first = self.builder.build_int_compare(
            IntPredicate::EQ,
            copy_idx,
            i64_type.const_int(0, false),
            "is_first",
        )?;

        // If first, use strcpy; otherwise, use strcat
        let first_block = self
            .context
            .append_basic_block(current_func, "join_first_str");
        let not_first_block = self
            .context
            .append_basic_block(current_func, "join_not_first_str");
        let after_copy_block = self
            .context
            .append_basic_block(current_func, "join_after_copy");

        self.builder
            .build_conditional_branch(is_first, first_block, not_first_block)?;

        // First string: use strcpy
        self.builder.position_at_end(first_block);
        self.builder.build_call(
            strcpy_fn,
            &[result_str.into(), copy_str.into()],
            "strcpy_first",
        )?;
        self.builder.build_unconditional_branch(after_copy_block)?;

        // Not first string: append separator then string
        self.builder.position_at_end(not_first_block);
        self.builder.build_call(
            strcat_fn,
            &[result_str.into(), separator_str_ptr.into()],
            "strcat_sep",
        )?;
        self.builder.build_call(
            strcat_fn,
            &[result_str.into(), copy_str.into()],
            "strcat_str",
        )?;
        self.builder.build_unconditional_branch(after_copy_block)?;

        // After copy: increment and continue
        self.builder.position_at_end(after_copy_block);
        let next_copy_idx =
            self.builder
                .build_int_add(copy_idx, i64_type.const_int(1, false), "next_copy_idx")?;
        self.builder.build_store(idx_alloca, next_copy_idx)?;
        self.builder.build_unconditional_branch(copy_loop)?;

        // Return block
        self.builder.position_at_end(return_block);

        // Create PHI node for result
        let result_phi = self.builder.build_phi(i8_ptr_type, "join_result")?;

        // Empty string case: return ""
        let empty_str = self.builder.build_global_string_ptr("", "empty_str")?;
        result_phi.add_incoming(&[(&empty_str, empty_check_block)]);

        // Single string case: return first string
        result_phi.add_incoming(&[(&first_str, single_check_block)]);

        // Normal case: return concatenated result
        result_phi.add_incoming(&[(&result_str, copy_loop)]);

        Ok(result_phi.as_basic_value())
    }


    /// Generate string.repeat(n) method call
    /// Repeats the string n times and returns the result
    ///
    /// Example: "abc".repeat(3) returns "abcabcabc"
    pub(in crate::llvm_backend) fn generate_string_repeat_method(
        &mut self,
        var_name: &str,
        arguments: &[Expression],
    ) -> Result<BasicValueEnum<'ctx>> {
        // Validate arguments
        if arguments.len() != 1 {
            return Err(CompilerError::internal(
                "repeat() requires exactly one argument (count: int)",
            ));
        }

        // Get the string pointer
        let string_ptr = if let Some(var_value) = self.variable_map.get(var_name) {
            var_value.into_pointer_value()
        } else {
            return Err(CompilerError::internal(&format!(
                "Undefined variable: {}",
                var_name
            )));
        };

        // Load the string pointer value (i8*)
        let i8_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
        let string_value = self
            .builder
            .build_load(i8_ptr_type, string_ptr, "string_val")?;
        let original_str_ptr = string_value.into_pointer_value();

        // Generate the count argument
        let count_value = self.generate_expression(&arguments[0])?;
        let mut count_int = count_value.into_int_value();

        // Convert to i64 if necessary (use sign-extend for negative numbers)
        if count_int.get_type() != self.context.i64_type() {
            count_int =
                self.builder
                    .build_int_s_extend(count_int, self.context.i64_type(), "count_i64")?;
        }

        let i8_type = self.context.i8_type();
        let i32_type = self.context.i32_type();
        let i64_type = self.context.i64_type();

        // Get strlen function
        let strlen_fn_type = i32_type.fn_type(&[i8_ptr_type.into()], false);
        let strlen_fn = if let Some(func) = self.module.get_function("strlen") {
            func
        } else {
            self.module.add_function("strlen", strlen_fn_type, None)
        };

        // Create basic blocks
        let current_func = self.current_function.unwrap();
        let zero_or_neg_check = self
            .context
            .append_basic_block(current_func, "repeat_zero_or_neg_check");
        let one_check = self
            .context
            .append_basic_block(current_func, "repeat_one_check");
        let calc_len_block = self
            .context
            .append_basic_block(current_func, "repeat_calc_len");
        let alloc_block = self
            .context
            .append_basic_block(current_func, "repeat_alloc");
        let copy_loop = self
            .context
            .append_basic_block(current_func, "repeat_copy_loop");
        let copy_body = self
            .context
            .append_basic_block(current_func, "repeat_copy_body");
        let return_block = self
            .context
            .append_basic_block(current_func, "repeat_return");

        // Jump to zero/negative check
        self.builder.build_unconditional_branch(zero_or_neg_check)?;

        // Check if count <= 0
        self.builder.position_at_end(zero_or_neg_check);
        let is_zero_or_neg = self.builder.build_int_compare(
            IntPredicate::SLE,
            count_int,
            i64_type.const_int(0, false),
            "is_zero_or_neg",
        )?;
        self.builder
            .build_conditional_branch(is_zero_or_neg, return_block, one_check)?;

        // Check if count == 1
        self.builder.position_at_end(one_check);
        let is_one = self.builder.build_int_compare(
            IntPredicate::EQ,
            count_int,
            i64_type.const_int(1, false),
            "is_one",
        )?;
        self.builder
            .build_conditional_branch(is_one, return_block, calc_len_block)?;

        // Calculate original string length
        self.builder.position_at_end(calc_len_block);
        let orig_len_i32 = self
            .builder
            .build_call(strlen_fn, &[original_str_ptr.into()], "orig_len_i32")?
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_int_value();
        let orig_len = self
            .builder
            .build_int_z_extend(orig_len_i32, i64_type, "orig_len")?;

        self.builder.build_unconditional_branch(alloc_block)?;

        // Allocate memory: orig_len * count + 1
        self.builder.position_at_end(alloc_block);
        let total_len = self
            .builder
            .build_int_mul(orig_len, count_int, "total_len")?;
        let total_with_null = self.builder.build_int_add(
            total_len,
            i64_type.const_int(1, false),
            "total_with_null",
        )?;

        // Allocate memory using malloc
        let malloc_fn_type = i8_ptr_type.fn_type(&[i64_type.into()], false);
        let malloc_fn = if let Some(func) = self.module.get_function("malloc") {
            func
        } else {
            self.module.add_function("malloc", malloc_fn_type, None)
        };

        let result_str = self
            .builder
            .build_call(malloc_fn, &[total_with_null.into()], "result_str")?
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // Get strcpy and strcat functions
        let strcpy_fn_type = i8_ptr_type.fn_type(&[i8_ptr_type.into(), i8_ptr_type.into()], false);
        let strcpy_fn = if let Some(func) = self.module.get_function("strcpy") {
            func
        } else {
            self.module.add_function("strcpy", strcpy_fn_type, None)
        };

        let strcat_fn_type = i8_ptr_type.fn_type(&[i8_ptr_type.into(), i8_ptr_type.into()], false);
        let strcat_fn = if let Some(func) = self.module.get_function("strcat") {
            func
        } else {
            self.module.add_function("strcat", strcat_fn_type, None)
        };

        // Initialize result with empty string
        self.builder
            .build_store(result_str, i8_type.const_int(0, false))?;

        // Create loop counter
        let idx_alloca = self.builder.build_alloca(i64_type, "idx")?;
        self.builder
            .build_store(idx_alloca, i64_type.const_int(0, false))?;

        self.builder.build_unconditional_branch(copy_loop)?;

        // Copy loop: repeat the string count times
        self.builder.position_at_end(copy_loop);
        let idx = self
            .builder
            .build_load(i64_type, idx_alloca, "idx")?
            .into_int_value();

        // Check if idx < count
        let continue_copy = self.builder.build_int_compare(
            IntPredicate::SLT,
            idx,
            count_int,
            "continue_copy",
        )?;
        self.builder
            .build_conditional_branch(continue_copy, copy_body, return_block)?;

        // Copy body: append the original string
        self.builder.position_at_end(copy_body);

        // Check if this is the first iteration
        let is_first = self.builder.build_int_compare(
            IntPredicate::EQ,
            idx,
            i64_type.const_int(0, false),
            "is_first",
        )?;

        let first_block = self
            .context
            .append_basic_block(current_func, "repeat_first");
        let not_first_block = self
            .context
            .append_basic_block(current_func, "repeat_not_first");
        let after_copy_block = self
            .context
            .append_basic_block(current_func, "repeat_after_copy");

        self.builder
            .build_conditional_branch(is_first, first_block, not_first_block)?;

        // First iteration: use strcpy
        self.builder.position_at_end(first_block);
        self.builder.build_call(
            strcpy_fn,
            &[result_str.into(), original_str_ptr.into()],
            "strcpy_first",
        )?;
        self.builder.build_unconditional_branch(after_copy_block)?;

        // Not first: use strcat
        self.builder.position_at_end(not_first_block);
        self.builder.build_call(
            strcat_fn,
            &[result_str.into(), original_str_ptr.into()],
            "strcat",
        )?;
        self.builder.build_unconditional_branch(after_copy_block)?;

        // After copy: increment index
        self.builder.position_at_end(after_copy_block);
        let next_idx = self
            .builder
            .build_int_add(idx, i64_type.const_int(1, false), "next_idx")?;
        self.builder.build_store(idx_alloca, next_idx)?;
        self.builder.build_unconditional_branch(copy_loop)?;

        // Return block with PHI node
        self.builder.position_at_end(return_block);
        let result_phi = self.builder.build_phi(i8_ptr_type, "repeat_result")?;

        // Empty string case (count <= 0)
        let empty_str = self.builder.build_global_string_ptr("", "empty_str")?;
        result_phi.add_incoming(&[(&empty_str, zero_or_neg_check)]);

        // Single copy case (count == 1): return original string
        result_phi.add_incoming(&[(&original_str_ptr, one_check)]);

        // Multiple copies case: return result
        result_phi.add_incoming(&[(&result_str, copy_loop)]);

        Ok(result_phi.as_basic_value())
    }
}
